import { Response } from 'express';

type SSEConnection = {
  res: Response;
  userId: string;
};

export class SseManager {
  private connections: Map<string, SSEConnection[]> = new Map();

  register(userId: string, res: Response) {
    const conn: SSEConnection = { res, userId };
    if (!this.connections.has(userId)) {
      this.connections.set(userId, []);
    }
    this.connections.get(userId)!.push(conn);

    res.setHeader('Content-Type', 'text/event-stream');
    res.setHeader('Cache-Control', 'no-cache');
    res.setHeader('Connection', 'keep-alive');
    res.flushHeaders();

    res.write(`data: ${JSON.stringify({ message: 'SSE connected' })}\n\n`);

    res.on('close', () => {
      const list = this.connections.get(userId);
      if (list) {
        this.connections.set(
          userId,
          list.filter(c => c.res !== res)
        );
      }
    });
  }

  sendTo(userId: string, data: any) {
    const connections = this.connections.get(userId);
    if (!connections) return;
    const msg = `data: ${JSON.stringify(data)}\n\n`;
    for (const conn of connections) {
      conn.res.write(msg);
    }
  }


  close(userId: string): boolean {
    const conns = this.connections.get(userId);
    if (!conns || conns.length === 0) return false;

    for (const conn of conns) {
      conn.res.end(); // 关闭连接
    }

    this.connections.delete(userId);
    return true;
  }

}

export const sseManager = new SseManager();
